#!/usr/bin/env python3

"""
Given n non-negative integers representing an elevation map where the width of each bar is 1, 
compute how much water it is able to trap after raining.

Example:
Input: [0, 1, 0, 2, 1, 0, 1, 3, 2, 1, 2, 1]
Output: 6
"""

class Solution:
    def trap(self, height):
        start_index = None
        for index, val in enumerate(height):
            if val == 0:
                continue
            start_index = index
            break
        if start_index is None:
            return 0
        length = len(height)
        result = 0
        while start_index < length:
            current = height[start_index]
            water_height, end_index = self.water_height_and_end_index(height, start_index + 1, length, current)
            if water_height == 0:
                start_index += 1
            else:
                result += self.area(height, start_index, end_index, water_height)
                start_index = end_index
        return result
            

    def area(self, height, start_index, end_index, water_height):
        result = 0
        while start_index < end_index - 1:
            result += water_height - height[start_index + 1]
            start_index += 1
        return result

    def water_height_and_end_index(self, height, start, length, target):
        max_height = 0
        max_index = None
        while start < length:
            current = height[start]
            if current >= target:
                max_height = current
                max_index = start
                break
            if current > max_height:
                max_height = current
                max_index = start
            start += 1
        max_height = min(target, max_height)
        return max_height, max_index
        

if __name__ == '__main__':
    solution = Solution()
    assert 0 == solution.trap([])
    input = [0,1,0,2,1,0,1,3,2,1,2,1]
    assert 6 == solution.trap(input)
